
from libatomism import *
from math import *
import array 
import matplotlib.pyplot as plt

# Set the logger to print only error messages
msLogger.setPriorityMin(Priority.ERROR)

# Create a units manager object, it will be used to instantiate 
# to define the unit system of the objects. It defines the unit 
# of the intput and output by default
unit=msUnitsManager.New("Angstrom Degree amu kcal/mol ps")

# Create the H2 ZMatrix:
# H
# H 0 0.7414
zmat=msZMat.New(unit).setId("zmat").set(" H H 0 0.7414")
print zmat

# Create the system: add the zmat entity 
system=msSystem.New(unit).addEntity(zmat)
print system

# Create a linear rigid rotor from the system
rotor = msLinearRigidRotor.New(unit).setMmtOfInertia(system).setSymmetryNumber(2)
print rotor

# Define the Emax parameter, look at the previous output (print rotor) to 
# understand its signification
rotor.getParameters().setParameter("Emax",150,unit)

# Create the translational motion
trans = msTranslation.New(unit).setSystem(system)

# Create the vibration motion
# step 1: defines a variable q0 in [0.1,5] Angstrom with current value at 0.7414 Angstrom, 
# and with an integration and derivative steps of 0.1 Angstrom
q0 = msScalarVariable.New("Angstrom").set(0.7414, 0.1, 5, 0.1, 0.1).setId("q0")

# step 2: defines the generalized coordinates; here only q0
gcoors = msGeneralizedCoordinates.New(unit).addVar(q0)
print gcoors

# step 3: creates the scalar function f(q0)=q0, it will be used to define the kinetic 
# operator.
kinfct = msParser.New("Angstrom").setCoordinates(gcoors).setExpression("q0").setId("f=q0")

# step 4: defines the kinetic operator, it describes how the masses move with respect 
# to the generalized coordinates. Firts we set the generalized coordinates and
# the system, then the degrees of freedom number 0 of 'zmat' (i.e. separation H-H) 
# is defined as the function 'kinfct'
kinop = msKineticOperator.New(unit).set(gcoors,system).addDynamicDof(zmat,0,kinfct)
print kinop

# step 5: defines the potential energy: a Morse potential with De=100 kcal/mol
De=100
epot = msParser.New("kcal/mol").setCoordinates(gcoors).setExpression("De * ( 1 - exp( - a*( q0 - re ) ))^2").setId("f=Morse")
# set the constants of 'epot'
epot.setConstante("De",De).setConstante("a",1.5).setConstante("re",0.7414)
print epot

# step 6: create a one dimensional quantum motion by setting the potential energy
# and the kinetic operator previously defined. Also set the parameters 'BasisSize' and 'Emax' (print qmotion for
# more information about them)
qmotion = msQuantumMotion1d.New(unit).setEqOfMotion(kinop,epot)
qmotion.getParameters().setParameter("BasisSize",501,unit)
qmotion.getParameters().setParameter("Emax",De,unit)
print qmotion

# we want to compare our exact treatment to an harmonic oscilator model
# we create an assembly of oscilator composed by one harmonic oscilator
oscillsHO = msOscillatorsAssembly.New(unit)
oscillsHO.addOscillator( msHarmonicOscillator.New("cm-1").set(3200.00) )

# create 2*2 plots
fig, axes = plt.subplots(2, 2)

# compute the eigenvalues
qmotion.computeEigenValues()

# get the 0, 5, 10 eigenvector
ev0  = qmotion.eigenvector(0)
ev5  = qmotion.eigenvector(5)
ev10 = qmotion.eigenvector(10)

# on the first subplot plots the potential energy and the eigenvectors with a sampling of 500 along
# the first coordinates (indiced 0, i.e. q0); the '200' is a scaling factor for having a nice drawing
# of the eigenvectors
axes[0, 0].plot(    epot.getXSample1d(0,500), epot.getYSample1d(0,500), "-"
		    ev0.getXSample1d(0,500) , [ 200*x + qmotion.eigenvalue(0)  for x in ev0.getYSample1d(0,500)] ,"-",
		    ev5.getXSample1d(0,500) , [ 200*x + qmotion.eigenvalue(5)  for x in ev5.getYSample1d(0,500)] ,"-",
		    ev10.getXSample1d(0,500), [ 200*x + qmotion.eigenvalue(10) for x in ev10.getYSample1d(0,500)],"-"
		)

# On the second subplot we plot the density of states and the sum of states of 'qmotion' and compare them 
# with those of the oscillator
axes[0, 1].plot(    qmotion.DOS().getXSample1d(0,500),    qmotion.DOS().getYSample1d(0,500),   "-"
		    qmotion.SOS().getXSample1d(0,500),    qmotion.SOS().getYSample1d(0,500),   "-"
		    oscillsHO.DOS().getXSample1d(0,500),  oscillsHO.DOS().getYSample1d(0,500), "-"
		    oscillsHO.SOS().getXSample1d(0,500),  oscillsHO.SOS().getYSample1d(0,500), "-"
		)
		
#this is the value of the kinetic function at the current coordinate:
print  kinop.kineticFunction()
# because it is a constant in this case we freeze its value for 'qmotion'
qmotion.freezeKinFct( kinop.kineticFunction() )

# compute the entropy and calorific capacity lists and plot them in a subplots. Again,
# the results are compared to those of the harmonic oscillator
s=[]
sHO=[]
for T in Tlist: 
    s.append( trans.S(T)  + rotor.S(T) + qmotion.S(T)  + unit.convert("kcal/mol",R) )
    sHO.append( trans.S(T)  + rotor.S(T) + oscillsHO.S(T)  + unit.convert("kcal/mol",R) )

cp=[]
cpHO=[]
for T in Tlist:
    cp.append( trans.Cv(T)+ rotor.Cv(T)+ qmotion.Cv(T) +  unit.convert("kcal/mol",R) )
    cpHO.append( trans.Cv(T)+ rotor.Cv(T)+ oscillsHO.Cv(T) +  unit.convert("kcal/mol",R) )

axes[0, 1].plot( Tlist, s ,    "-", Tlist, sHO)
axes[1, 1].plot( Tlist, cp[:], "-", Tlist, cpHO)

# this is the experimental values, coming from the janaf table (http://kinetics.nist.gov/janaf/html/H-050.html),
# defined in J/mol
TListexp  = [200,300,500,800,1200,1500,2000,2500,3000] 
ExpList   = [119.412,130.858,145.737,159.548,171.79,178.846,188.418,196.243,202.891]   
ExpListCp = [27.447,28.849,29.260,29.624,30.992,32.298,34.280,35.842,37.087] 
 
axes[0, 1].plot( TListexp, [ unit.convert("J/mol",x) for x in ExpList],   "+" )
axes[1, 1].plot( TListexp, [ unit.convert("J/mol",x) for x in ExpListCp], "+" )


fig.show()


 
